package hk.polyu;

import java.io.IOException;
import java.io.StreamTokenizer;
import java.io.StringReader;
import java.util.*;


/**
 * This class is used to search specific query using 
 * the prebuilt index files.
 * @author hezongjian
 *
 */
public class Searcher {
	
	/**
	 * The inverted index to be used for search.
	 */
	InvertedIndex 	mIndex;
	
	/**
	 * The document index to be used for search.
	 */
	DocumentIndex	mDocIndex;
	
	/**
	 * The only constructor accepting two parameters for the index
	 * @param i Inverted Index
	 * @param d Document Index
	 */
	public Searcher(InvertedIndex i, DocumentIndex d) {
		this.mDocIndex = d;
		this.mIndex = i;
	}
	
	/**
	 * This is used to sort QueryResult according to document ID
	 */
	static final Comparator<QueryResult> ID_ORDER = new Comparator<QueryResult>() {
		@Override
		public int compare(QueryResult o1, QueryResult o2) {
			return o1.mID - o2.mID;
		}
	};
	
	/**
	 * This is used to sort QueryResult in desc order according to weight
	 */
	static final Comparator<QueryResult> WEIGHT_ORDER_DESC = new Comparator<QueryResult>() {
		@Override
		public int compare(QueryResult o1, QueryResult o2) {
			return Double.compare(o2.mWeight, o1.mWeight);
		}
	};
	
	
	/**
	 * This method search the input terms with index and return a list of result.
	 * The vector space model is used to search the result.
	 * The result is sorted according to the document weight in descending order.
	 * @param terms The terms that will be searched. If multiple terms are specified. OR operator are used.
	 * @return The sorted query result.
	 */
	public Vector<QueryResult> queryOR(Vector<String> terms) {
		// weight matrix. one row per term.
		int docN = this.mDocIndex.getDocumentCount();
		int termN = terms.size();
		double [][] w = new double[termN][docN];
		
		for (int i = 0; i < termN; i++) {
			String s = terms.elementAt(i);
			// Firstly, stem it.
			Stemmer stm = new Stemmer();
			char [] buf = s.toCharArray();
			stm.add(buf, buf.length);
			stm.stem();
			String stemmedTerm = stm.toString();
			
			// Calculate the TF-IDF for each document.
			for (int j = 0; j < docN; j++) {
				w[i][j] = this.getTfIdf(stemmedTerm, j);
			}
		}
		// get the weighted document
		double []fw = new double[docN];
		for (int i = 0; i < docN; i++) {
			for (int j = 0; j < termN; j++) {
				fw[i] += w[j][i];
			}
		}
		// generate result
		Vector<QueryResult> qr = new Vector<QueryResult>();
		for(int i = 0; i < docN; i++) {
			if (fw[i] > 0.1) { 	// add with threshold.
				QueryResult r = new QueryResult(this.mDocIndex.mDocs.get(i), fw[i]);
				qr.add(r);
			}
		}
		// sorting.
		Collections.sort(qr, WEIGHT_ORDER_DESC);
		return qr;
	}
	
	
	/**
	 * The phase version of query. Only support one phrase.
	 * 
	 * @param terms
	 *            the phrase.
	 * @return The query containing the phrase.
	 */
	public Vector<QueryResult> queryPhrase(Vector<String> phrase) {
		// weight matrix. one row per term.
		int docN = this.mDocIndex.getDocumentCount();
		int phraseN = phrase.size();
		
		// stemming
		Vector<String> stemmedPhrase = new Vector<String>();
		for (int i = 0; i < phraseN; i++) {
			String s = phrase.elementAt(i);
			// Firstly, stem it.
			Stemmer stm = new Stemmer();
			char[] buf = s.toCharArray();
			stm.add(buf, buf.length);
			stm.stem();
			stemmedPhrase.add(stm.toString());
		}
		// generate result
		Vector<QueryResult> qr = new Vector<QueryResult>();
		for (int i = 0; i < docN; i++) {
			double tfidf = this.getTfIdf(stemmedPhrase, i);
			if ( tfidf > 0.1) { // add with threshold.
				QueryResult r = new QueryResult(this.mDocIndex.mDocs.get(i), tfidf);
				qr.add(r);
			}
		}
		// sorting.
		Collections.sort(qr, WEIGHT_ORDER_DESC);
		return qr;
	}
	
	
	/**
	 * Get the tf-idf weight for specific term in specific document.
	 * @param term The term which will be calculated
	 * @param docID The document ID.
	 * @return The tf-idf weight.
	 */
	public double getTfIdf(String term, int docID) {
		double tfidf = 0.0;
		// get TF
		double tf = 0.0;
		double idf = 0.0;
		int hasTerm = 0;
		HashMap<Integer, Vector<Integer>> inner = this.mIndex.mTokens.get(term);
		if (inner != null) {
			hasTerm = inner.keySet().size();
			Vector<Integer> docs = inner.get(docID);
			if (docs != null) {
				tf = docs.size();
			}
		}
		if (tf == 0.0) {
			return tfidf;
		}
		
		// get IDF
		int docN = this.mDocIndex.getDocumentCount();
		idf = (double) Math.log( (double)docN/ (hasTerm+1.0));
		
		// calculate and return
		tfidf = tf * idf;
		return tfidf;
	}
	
	/**
	 * This method checks whether a document contain specific phrase.
	 * @param phrase
	 * @param docID
	 * @return
	 */
	public Vector<Integer> getPhrase(Vector<String> phrase, int docID) {
		Vector<Integer> ret = new Vector<Integer>();
		@SuppressWarnings("unchecked")
		Vector<Integer> [] vecs = new Vector[3];
		
		for(int i = 0; i < phrase.size(); i++) {
			String term = phrase.get(i);
			HashMap<Integer, Vector<Integer>> inner = this.mIndex.mTokens.get(term);
			if (inner == null || ! inner.containsKey(docID))
				return null;
			vecs[i] = inner.get(docID);
		}
		
		for (int i = 0; i < vecs[0].size(); i++) {
			int iPos = vecs[0].get(i);
			int j = 1;
			while (j < phrase.size()) {
				if (vecs[j].contains(iPos + j) )
					j++;
				else
					break;
			}
			if (j == phrase.size())
				ret.add(iPos);
		}
		
		if (ret.size() == 0)
			ret = null;
		return ret;
	}
	
	/**
	 * The phrase version of tf-idf calculation.
	 * @param term
	 * @param docID
	 * @return
	 */
	public double getTfIdf(Vector<String> phrase, int docID) {
		double tfidf = 0.0;
		double tf = 0.0;
		double idf = 0.0;
		
		// get TF
		Vector<Integer> pos = this.getPhrase(phrase, docID);
		if (pos != null) {
			tf = pos.size();
		}

		if (tf == 0.0) {
			return tfidf;
		}
		
		// get IDF
		int docN = this.mDocIndex.getDocumentCount();
		int hasTerm = 0;
		for(int i = 0; i < docN; i++){
			if (this.getPhrase(phrase, i) != null)
				hasTerm++;
		}
		idf = Math.log( (double)docN/ (hasTerm+1.0));

		// calculate and return
		tfidf = tf * idf;
		return tfidf;		
	}
	
	/**
	 * Calculate the intersection of two query results.
	 * @param a One query results
	 * @param b The other query result
	 * @return The intersected query result.
	 */
	public static Vector<QueryResult> andResults(Vector<QueryResult> a, Vector<QueryResult> b) {
		Vector<QueryResult> ret = new Vector<QueryResult>();		
		Vector<QueryResult> copyA = new Vector<QueryResult>(a);
		Vector<QueryResult> copyB = new Vector<QueryResult>(b);
		
		// sort two input
		Collections.sort(copyA, ID_ORDER);
		Collections.sort(copyB, ID_ORDER);
		int i = 0, j = 0;
		
		while (i < copyA.size() && j < copyB.size()) {
			QueryResult ta = copyA.get(i);
			QueryResult tb = copyB.get(j);
			if(ta.mID == tb.mID) {
				if (ta.mWeight > tb.mWeight)
					ret.add(ta);
				else
					ret.add(tb);
				i++;
				j++;
			}
			else if (ta.mID < tb.mID) {
				i++;
			}
			else {
				j++;
			}
		} // end while		
		Collections.sort(ret, WEIGHT_ORDER_DESC);
		return ret;
	}

	/**
	 * This method find out the OR of two query results sets.
	 * @param a result A
	 * @param b result B
	 * @return The combination of the two results.
	 */
	public static Vector<QueryResult> orResults(Vector<QueryResult> a, Vector<QueryResult> b) {
		
		if (a.size() == 0)
			return b;
		
		if(b.size() == 0)
			return a;
		
		Vector<QueryResult> ret = new Vector<QueryResult>();		
		Vector<QueryResult> copyA = new Vector<QueryResult>(a);
		Vector<QueryResult> copyB = new Vector<QueryResult>(b);
		
		// sort two input
		Collections.sort(copyA, ID_ORDER);
		Collections.sort(copyB, ID_ORDER);
		int i = 0, j = 0;
		
		while (i < copyA.size() && j < copyB.size()) {
			QueryResult ta = copyA.get(i);
			QueryResult tb = copyB.get(j);
			if(ta.mID == tb.mID) {
				ta.mWeight += tb.mWeight;
				ret.add(ta);
				i++;
				j++;
			}
			else if (ta.mID < tb.mID) {
				ret.add(ta);
				i++;
			}
			else {
				ret.add(tb);
				j++;
			}
		} // end while
		
		// add the remaining
		while(i < copyA.size()) {
			ret.add(copyA.get(i));
			i++;
		}
		while(j < copyB.size()) {
			ret.add(copyB.get(j));
			j++;
		}

		// sort it
		Collections.sort(ret, WEIGHT_ORDER_DESC);
		return ret;
	}
	
	
	private final String s1="AND";
	private final String s2="OR";
	
	public Vector<QueryResult> query(String q) {
		Vector<QueryResult> qr = null;
		Stack<String> stack_operand = new Stack<String>();
		Stack<String> stack_operator = new Stack<String>();
		StreamTokenizer tokens = new StreamTokenizer(new StringReader(q));
		
		tokens.quoteChar('\"');
		tokens.quoteChar('\'');

		// push all tokens into stack
		try {
			while(tokens.nextToken()!=StreamTokenizer.TT_EOF) {
				String token = tokens.sval;

				if (!s1.equals(token) && !s2.equals(token))
					stack_operand.push(token);
				else 
					stack_operator.push(token);
			}
		} catch (IOException e) {
			e.printStackTrace();
		}
		
		// Evaluate the operation.
		qr = this.queryOneToken(stack_operand.pop());
		while (! stack_operator.isEmpty()) {
			String oprtr = stack_operator.pop();
			Vector<QueryResult> cur = this.queryOneToken(stack_operand.pop());
			
			if(oprtr.equalsIgnoreCase(s1))
				qr = Searcher.andResults(qr, cur);
			else
				qr = Searcher.orResults(qr, cur);
		}
	
		return qr;
	}
	
	Vector<QueryResult> queryOneToken(String t) {
		String lt = t.toLowerCase();
		Vector<String> sv = new Vector<String>();
		sv.addAll(Arrays.asList(lt.split(" ")));
		
		if (sv.size() == 1) 
			return this.queryOR(sv);
		else
			return this.queryPhrase(sv);
	}
	
}
